package avl;

public class AvlTree {
    static class TreeNode {
        public int val;
        public TreeNode parent;
        public TreeNode left;
        public TreeNode right;
        //平衡因子
        public int bf;
        public TreeNode(int val) {
            this.val = val;
        }
    }
    //根结点
    public TreeNode root;
    //插入
    public void insert(int val) {
        TreeNode node = new TreeNode(val);
        if(root == null) {
            root = node;
            return;
        }
        TreeNode cur = root;
        TreeNode parent = root;
        while(cur != null) {
            parent = cur;
            if(cur.val > val) {
                cur = cur.left;
            }else if(cur.val < val) {
                cur = cur.right;
            }else {
                //cur.val == node.val
                return;
            }
        }
        if(parent.val > val) {
            parent.left = node;
        }else if(parent.val < val) {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        //插入完毕，就可以进行平衡的调整
        while(parent != null) {
            if(cur == parent.right) {
                parent.bf++;
            }else {
                parent.bf--;
            }

            if(parent.bf == 0) {
                //平衡因子以及调整完毕
                break;
            }else if(parent.bf == 1 || parent.bf == -1) {
                cur = parent;
                parent = parent.parent;
            }else {
                //prev.bf == 2 || prev.bf == -2
                //此时可以进行调整了
                if(parent.bf == -2) {
                    if(cur.bf == -1) {
                        //右单旋
                        rotateRight(parent);
                    }else {
                        //左右双旋
                        rotateLR(parent);
                    }
                }else {
                    if(cur.bf == 1) {
                        //左单旋(包含cur.bf == 0)
                        rotateLeft(parent);
                    }else {
                        //右左双旋
                        rotateRL(parent);
                    }
                }
                //上述步骤执行完毕之后就可以跳出循环了
                break;
            }
        }
    }
    //判断当前是否是AVL树（高度平衡的二叉搜索树）
    private int height(TreeNode root) {
        if(root == null) {
            return 0;
        }
        int leftHeight = height(root.left);
        int rightHeight = height(root.right);

        int maxHeight = Math.max(leftHeight,rightHeight);
        return maxHeight + 1;
    }
    public boolean isBalance(TreeNode root) {
        if(root == null) {
            return true;
        }
        int leftHeight = height(root.left);
        int rightHeight = height(root.right);
        if(rightHeight-leftHeight != root.bf) {
            System.out.println("当前节点的平衡因子出现异常 bf = " + root.bf);
            return false;
        }
        return Math.abs(leftHeight-rightHeight) <= 1
                && isBalance(root.left)
                && isBalance(root.right);
    }

    private void rotateRL(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;
        rotateRight(subR);
        rotateLeft(parent);

        if(bf == 1) {
            parent.bf = -1;
            subR.bf = 0;
            subRL.bf = 0;
        }else if(bf == -1) {
            //这里必须要加bf == -1,否则会出现错误
            parent.bf = 0;
            subR.bf = 1;
            subRL.bf = 0;
        }
    }

    private void rotateLR(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;
        rotateLeft(subL);
        rotateRight(parent);
        if(bf == 1) {
            parent.bf = 0;
            subL.bf = -1;
            subLR.bf = 0;
        }else if(bf == -1){
            //这里必须要加bf == -1,否则会出现错误
            parent.bf = 1;
            subL.bf = 0;
            subLR.bf = 0;
        }//还有一种为0的情况，在上述进行左单旋和右单旋就已经进行了处理
    }

    private void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRl = subR.left;
        subR.left = parent;
        TreeNode pParent = parent.parent;
        parent.parent = subR;
        parent.right = subRl;
        if(subRl != null) {
            subRl.parent = parent;
        }
        if(root == parent) {
            root = subR;
            root.parent = null;
        }else {
            if(parent == pParent.left) {
                pParent.left = subR;
            }else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        subR.bf = 0;
        parent.bf = 0;
    }

    private void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;

        parent.left = subLR;
        subL.right = parent;
        TreeNode pParent = parent.parent;

        parent.parent = subL;
        if(subLR != null) {
            subLR.parent = parent;
        }

        if(root == parent) {
            //更新当前新的根节点
            root = subL;
            root.parent = null;
        }else {
            //不是根节点，判断当前pParent节点是左子树还是右子树
            if(parent == pParent.right) {
                pParent.right = subL;
            }else {
                pParent.left = subL;
            }
            subL.parent = pParent;
        }
        parent.bf = 0;
        subL.bf = 0;
    }
}
